package container

import (
	"com.yy.api/common/logger"
	"sync"
)

type GpuState struct {
	mu   sync.Mutex
	used bool
}
type GpuManager struct {
	mu   sync.Mutex
	gpus map[int]*GpuState
}

func NewGpuManager(total []int) *GpuManager {
	m := &GpuManager{
		gpus: make(map[int]*GpuState, len(total)),
	}
	for _, id := range total {
		m.gpus[id] = &GpuState{}
	}
	return m
}

func (g *GpuManager) Acquire(gpuIds []int) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	var acquired []int
	for _, id := range gpuIds {
		state, exists := g.gpus[id]
		if !exists {
			goto rollback
		}
		state.mu.Lock()
		if state.used {
			state.mu.Unlock()
			goto rollback
		}
		state.used = true
		state.mu.Unlock()
		acquired = append(acquired, id)
	}
	return true

rollback:
	for _, id := range acquired {
		if s, ok := g.gpus[id]; ok {
			s.mu.Lock()
			s.used = false
			s.mu.Unlock()
		}
	}
	return false
}

func (g *GpuManager) Release(gpuId int) bool {
	g.mu.Lock()
	defer g.mu.Unlock()
	state, exists := g.gpus[gpuId]
	if !exists {
		logger.Default.Errorf("gpu not exists! gpuId:%v", gpuId)
		return false
	}
	state.mu.Lock()
	defer state.mu.Unlock()

	if !state.used {
		logger.Default.Infof("Trying to release an unreleased GPU:%v", gpuId)
		return false
	}
	state.used = false
	logger.Default.Infof("Released GPU: %v", gpuId)
	return true
}

func (g *GpuManager) IsAvailable(gpuNum int) ([]int, bool) {
	g.mu.Lock()
	defer g.mu.Unlock()
	var available []int
	for id, state := range g.gpus {
		state.mu.Lock()
		if !state.used {
			available = append(available, id)
		}
		state.mu.Unlock()
	}
	if len(available) < gpuNum {
		logger.Default.Errorf("Not enough GPUs available. Need %d, but only %d available.", gpuNum, len(available))
		return []int{}, false
	}
	var acquired []int
	for i := 0; i < gpuNum; i++ {
		acquired = append(acquired, available[i])
	}
	return acquired, true
}
